name: Build AutoGPTQ Wheels with CUDA for Linux

on: workflow_dispatch

jobs:
  build_wheels:
    if: ${{ github.repository_owner == 'AutoGPTQ' }}
    name: Build wheels for ${{ matrix.os }} and Python ${{ matrix.python }} and CUDA ${{ matrix.cuda }}
    runs-on: ${{ matrix.os }}
    strategy:
      matrix:
        os: [ubuntu-20.04]
        pyver: ["3.8", "3.9", "3.10", "3.11"]
        cuda: ["11.8"]  # wheel for 12.1 are built in build_wheels_pypi.yml
    defaults:
      run:
        shell: bash
    env:
        CUDA_VERSION: ${{ matrix.cuda }}

    steps:
      - uses: actions/checkout@v3

      - name: Free disk space
        run: |
          # Go from 19G to 54G free disk space in 3min
          df -h
          sudo apt-get update
          sudo apt-get purge -y '^apache.*'
          sudo apt-get purge -y '^imagemagick.*'
          sudo apt-get purge -y '^dotnet.*'
          sudo apt-get purge -y '^aspnetcore.*'
          sudo apt-get purge -y 'php.*'
          sudo apt-get purge -y '^temurin.*'
          sudo apt-get purge -y '^mysql.*'
          sudo apt-get purge -y '^java.*'
          sudo apt-get purge -y '^openjdk.*'
          sudo apt-get purge -y microsoft-edge-stable google-cloud-cli azure-cli google-chrome-stable firefox powershell mono-devel
          df -h
          sudo apt-get autoremove -y >/dev/null 2>&1
          sudo apt-get clean
          df -h
          echo "https://github.com/actions/virtual-environments/issues/709"
          sudo rm -rf "$AGENT_TOOLSDIRECTORY"
          df -h
          echo "remove big /usr/local"
          sudo rm -rf "/usr/local/share/boost"
          sudo rm -rf /usr/local/lib/android >/dev/null 2>&1
          df -h
          echo "remove /usr/share leftovers"
          sudo rm -rf /usr/share/dotnet/sdk > /dev/null 2>&1
          sudo rm -rf /usr/share/dotnet/shared > /dev/null 2>&1
          sudo rm -rf /usr/share/swift > /dev/null 2>&1
          df -h
          echo "remove other leftovers"
          sudo rm -rf /var/lib/mysql > /dev/null 2>&1
          sudo rm -rf /home/runner/.dotnet > /dev/null 2>&1
          sudo rm -rf /home/runneradmin/.dotnet > /dev/null 2>&1
          sudo rm -rf /etc/skel/.dotnet > /dev/null 2>&1
          sudo rm -rf /usr/local/.ghcup > /dev/null 2>&1
          sudo rm -rf /usr/local/aws-cli > /dev/null 2>&1
          sudo rm -rf /usr/local/lib/node_modules > /dev/null 2>&1
          sudo rm -rf /usr/lib/heroku > /dev/null 2>&1
          sudo rm -rf /usr/local/share/chromium > /dev/null 2>&1
          df -h

      - uses: actions/setup-python@v3
        with:
          python-version: ${{ matrix.pyver }}

      - name: Setup Miniconda
        uses: conda-incubator/setup-miniconda@v2.2.0
        with:
          activate-environment: "build"
          python-version: ${{ matrix.pyver }}
          mamba-version: "*"
          use-mamba: false
          channels: conda-forge,defaults
          channel-priority: true
          add-pip-as-python-dependency: true
          auto-activate-base: false

      - name: Install Dependencies
        run: |
          conda install cuda-toolkit -c "nvidia/label/cuda-${CUDA_VERSION}.0"

          # Refer to https://pytorch.org/get-started/locally/
          python -m pip install torch --index-url https://download.pytorch.org/whl/cu118
          python -m pip install --upgrade build setuptools wheel ninja numpy gekko pandas

      - name: Check install
        run: |
          python -c "import torch; print('torch version:', torch.__version__)"

      - name: Build Wheel
        run: |
          # For some reason $CONDA_PREFIX is empty.
          export CUDA_HOME=/usr/share/miniconda
          export CUDA_PATH=/usr/share/miniconda
          export LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:${CONDA_PREFIX}/lib"

          export TORCH_CUDA_ARCH_LIST="6.0 6.1 7.0 7.5 8.0 8.6 8.9 9.0+PTX"

          echo "CUDA_PATH:"
          echo $CUDA_PATH

          echo "PYPI_RELEASE:"
          echo $PYPI_RELEASE

          python setup.py sdist bdist_wheel

      - uses: actions/upload-artifact@v3
        with:
          name: 'linux-cuda-wheels'
          path: ./dist/*.whl
